from typing import Any, Dict, Optional, Type, Union, overload, Literal
import torch.nn as nn

from continuallearning.interfaces import HookManagerInterface
from continuallearning.models.pefts.common.config.hook_manager_config import (
    HookManagerConfig,
)
from ._base_hook_manager import AbstractHookManager
from .hook_manager import UnifiedHookManager


__all__ = [
    "HookManagerFactory",
    "create_hook_manager",
    "AbstractHookManager",
    "UnifiedHookManager",
]


class HookManagerFactory:
    """
    Factory class for creating hook managers.

    This class provides a unified way to create hook managers with consistent
    configuration handling.

    Example:
        >>> factory = HookManagerFactory()
        >>> config = HookManagerConfig(
        ...     hook_factory=HookFactoryConfig(hook_name="lora"),
        ...     discovery=DiscoveryConfig(find_linear_layers=True)
        ... )
        >>> manager = factory.create_manager(model, config=config)
    """

    def create_manager(
        self,
        model: nn.Module,
        manager_type: Union[str, Type[HookManagerInterface]] = "unified",
        config: Optional[Union[Dict[str, Any], HookManagerConfig]] = None,
    ) -> HookManagerInterface:
        """
        Create a hook manager instance.

        This method supports creating hook managers either by:
        1. Specifying a registered manager type name (string)
        2. Providing a hook manager class directly

        Args:
            model: Model to apply hooks to
            manager_type: Type of hook manager to create (class or string name)
            config: Configuration for the hook manager (dict or HookManagerConfig)

        Returns:
            HookManagerInterface: Instantiated hook manager

        Raises:
            ValueError: If manager_type is invalid or not registered
            TypeError: If manager_type is not a string or HookManagerInterface subclass
        """
        # Get manager class
        if isinstance(manager_type, str):
            from continuallearning.registry import HOOK_MANAGER_REGISTRY

            manager_cls = HOOK_MANAGER_REGISTRY.get(manager_type)
            if manager_cls is None:
                available_managers = list(HOOK_MANAGER_REGISTRY.list_registered())
                raise ValueError(
                    f"Hook manager type '{manager_type}' not registered. "
                    f"Available managers: {available_managers}"
                )
        elif isinstance(manager_type, type) and issubclass(
            manager_type, HookManagerInterface
        ):
            manager_cls = manager_type
        else:
            raise TypeError(
                f"manager_type must be a string or HookManagerInterface subclass, "
                f"got {type(manager_type).__name__}"
            )

        # Create the manager with unified interface
        return manager_cls(model=model, config=config)


@overload
def create_hook_manager(
    model: nn.Module,
    config: Union[Dict[str, Any], HookManagerConfig],
    manager_type: Literal["unified"],
) -> UnifiedHookManager: ...


@overload
def create_hook_manager(
    model: nn.Module,
    config: Union[Dict[str, Any], HookManagerConfig],
    manager_type: Literal["UNIFIED"],
) -> UnifiedHookManager: ...


def create_hook_manager(
    model: nn.Module,
    config: Union[Dict[str, Any], HookManagerConfig],
    manager_type: str,
) -> HookManagerInterface:
    """
    Create a hook manager with the specified configuration.

    This convenience function simplifies the creation of hook managers by
    providing a clean interface for the most common use case.

    Args:
        model: Model to apply hooks to
        config: Configuration for the hook manager, can be:
                - None: Use default configuration
                - dict: Configuration dictionary that will be converted to HookManagerConfig
                - HookManagerConfig: Pre-built configuration object
        manager_type: Type of hook manager to create (default: "unified")

    Returns:
        HookManagerInterface: Instantiated and configured hook manager

    Examples:
        >>> # Create with default configuration
        >>> manager = create_hook_manager(model)

        >>> # Create with dictionary configuration
        >>> manager = create_hook_manager(
        ...     model,
        ...     config={
        ...         "hook_factory": {"hook_name": "lora", "hook_config": {"r": 8}},
        ...         "discovery": {"find_linear_layers": True}
        ...     }
        ... )

        >>> # Create with HookManagerConfig object
        >>> config = HookManagerConfig(
        ...     hook_factory=HookFactoryConfig(hook_name="prefix"),
        ...     discovery=DiscoveryConfig(find_transformer_layers=True)
        ... )
        >>> manager = create_hook_manager(model, config=config)

        >>> # Create with custom manager type
        >>> manager = create_hook_manager(model, manager_type="custom_manager")
    """
    factory = HookManagerFactory()
    return factory.create_manager(
        model=model,
        manager_type=manager_type,
        config=config,
    )
